import numpy as np
from enum import Enum

class SubChannelType(Enum):
    """子通道类型枚举"""
    CORNER = "corner"    # 角通道
    EDGE = "edge"       # 边通道
    CENTER = "center"   # 中心通道

class SubChannel:
    """子通道类，定义单个子通道的几何和物理特性"""
    def __init__(self, id, channel_type, geometry):
        # 基本标识
        self.id = id
        self.type = channel_type  # 统一使用type作为类型属性
        
        # 几何参数 - 根据通道类型设置不同的默认值
        if channel_type == SubChannelType.CORNER:
            self.area = geometry.get('area', 4.293e-5)  # m^2
            self.wetted_perimeter = geometry.get('wetted_perimeter', 0.01246)  # m
            self.hydraulic_diameter = geometry.get('hydraulic_diameter', 0.01377)  # m
            self.gap_width = geometry.get('gap_width', 0.0033)  # m
            self.gap_length = geometry.get('gap_length', 0.0126)  # m
            self.centroid_distance = geometry.get('centroid_distance', 0.0167)  # m
            
        elif channel_type == SubChannelType.EDGE:
            self.area = geometry.get('area', 7.447e-5)  # m^2
            self.wetted_perimeter = geometry.get('wetted_perimeter', 0.02734)  # m
            self.hydraulic_diameter = geometry.get('hydraulic_diameter', 0.01089)  # m
            self.gap_width = geometry.get('gap_width', 0.0033)  # m
            self.gap_length = geometry.get('gap_length', 0.0126)  # m
            self.centroid_distance = geometry.get('centroid_distance', 0.0167)  # m
            
        else:  # CENTER
            self.area = geometry.get('area', 1.5876e-4)  # m^2
            self.wetted_perimeter = geometry.get('wetted_perimeter', 0.02985)  # m
            self.hydraulic_diameter = geometry.get('hydraulic_diameter', 0.02125)  # m
            self.gap_width = geometry.get('gap_width', 0.0033)  # m
            self.gap_length = geometry.get('gap_length', 0.0126)  # m
            self.centroid_distance = geometry.get('centroid_distance', 0.0167)  # m
        
        # 相邻子通道连接信息
        self.connections = []  # 存储相邻子通道的ID和连接类型
        
        # 场变量 - 两相流体
        self.alpha = np.zeros(1)  # 空泡份额
        self.pressure = np.zeros(1)  # 压力
        self.temperature = {
            'liquid': np.zeros(1),
            'vapor': np.zeros(1)
        }
        self.velocity = {
            'liquid': {
                'axial': np.zeros(1),    # 轴向速度
                'lateral': np.zeros(1)    # 横向速度
            },
            'vapor': {
                'axial': np.zeros(1),
                'lateral': np.zeros(1)
            }
        }
        self.density = {
            'liquid': np.zeros(1),
            'vapor': np.zeros(1)
        }
        
        # 存储上一时间步的值
        self.previous_alpha = np.zeros(1)
        self.previous_pressure = np.zeros(1)
        self.previous_velocity = {
            'liquid': {
                'axial': np.zeros(1),
                'lateral': np.zeros(1)
            },
            'vapor': {
                'axial': np.zeros(1),
                'lateral': np.zeros(1)
            }
        }
    
    def get_geometry_parameters(self):
        """获取子通道的几何参数
        Returns:
            dict: 包含所有几何参数的字典
        """
        return {
            'type': self.type,
            'area': self.area,
            'wetted_perimeter': self.wetted_perimeter,
            'hydraulic_diameter': self.hydraulic_diameter,
            'gap_width': self.gap_width,
            'gap_length': self.gap_length,
            'centroid_distance': self.centroid_distance
        }
    
    def get_connection_parameters(self, adjacent_channel):
        """获取与相邻子通道的连接参数
        Args:
            adjacent_channel: 相邻子通道对象
        Returns:
            dict: 包含连接参数的字典
        """
        # 计算两个子通道中心的距离
        centroid_distance = self.centroid_distance
        
        # 获取连接类型（轴向或横向）
        connection_type = 'lateral'  # 默认为横向连接
        for conn in self.connections:
            if conn['neighbor_id'] == adjacent_channel.id:
                connection_type = conn['type']
                break
        
        return {
            'connection_type': connection_type,
            'gap_width': self.gap_width,
            'gap_length': self.gap_length,
            'centroid_distance': centroid_distance
        }
    
    def add_connection(self, neighbor_id, connection_type, gap_width):
        """添加与相邻子通道的连接信息"""
        self.connections.append({
            'neighbor_id': neighbor_id,
            'type': connection_type,  # 'axial' 或 'lateral'
            'gap_width': gap_width
        })
        
    def save_current_state(self):
        """保存当前状态为上一时间步的值"""
        self.previous_alpha = np.copy(self.alpha)
        self.previous_pressure = np.copy(self.pressure)
        for phase in ['liquid', 'vapor']:
            for direction in ['axial', 'lateral']:
                self.previous_velocity[phase][direction] = np.copy(self.velocity[phase][direction])

class Grid:
    """网格类，管理所有子通道和其连接关系"""
    def __init__(self, n_axial=20, channel_length=3.66):
        self.n_axial = n_axial  # 轴向节点数
        self.n_channels = 36     # 6x6网格的子通道总数
        self.subchannels = []    # 存储所有子通道
        self.channel_length = channel_length  # 通道长度
        self.dz = channel_length / (n_axial - 1)  # 轴向网格间距
        self.dx = 0.0126  # 横向网格间距（根据实际燃料棒间距设置）
        
        # 时间步相关
        self.current_time = 0.0  # 当前时间
        self.dt = 0.0  # 当前时间步长
        self.time_step_number = 0  # 当前时间步编号
        
        # 网格索引映射
        self.axial_indices = np.arange(n_axial)
        self.channel_indices = np.arange(self.n_channels)
        
        # 子通道类型映射
        self.channel_types = self._initialize_channel_types()
        
        # 创建所有子通道
        self._create_all_subchannels()
        
        # 建立子通道之间的连接关系
        self.setup_connections()
        
    def _initialize_channel_types(self):
        """初始化6x6网格的子通道类型映射"""
        # 1代表角通道，2代表边通道，3代表中心通道
        channel_type_map = [
            [1, 2, 2, 2, 2, 1],
            [2, 3, 3, 3, 2, 2],
            [2, 3, 3, 3, 2, 2],
            [2, 3, 3, 3, 2, 2],
            [2, 2, 2, 2, 2, 2],
            [1, 2, 2, 2, 2, 1]
        ]
        channel_types = {}
        for i in range(6):
            for j in range(6):
                channel_id = i * 6 + j
                if channel_type_map[i][j] == 1:
                    channel_types[channel_id] = SubChannelType.CORNER
                elif channel_type_map[i][j] == 2:
                    channel_types[channel_id] = SubChannelType.EDGE
                else:
                    channel_types[channel_id] = SubChannelType.CENTER
        return channel_types
        
    def create_subchannel(self, id):
        """创建新的子通道，根据ID确定类型和几何参数"""
        channel_type = self.channel_types[id % 36]  # 使用取模运算处理多层
        geometry = {}  # 几何参数将在SubChannel类中根据类型设置默认值
        subchannel = SubChannel(id, channel_type, geometry)
        self.subchannels.append(subchannel)
        return subchannel
    
    def setup_connections(self):
        """建立子通道之间的连接关系"""
        for subchannel in self.subchannels:
            base_id = subchannel.id % 36  # 当前层内的ID
            row = base_id // 6
            col = base_id % 6
            
            # 轴向连接
            if subchannel.id < self.n_channels * (self.n_axial - 1):
                subchannel.add_connection(
                    subchannel.id + 36,  # 连接到上一层对应位置
                    'axial',
                    self.dz
                )
            
            # 横向连接
            # 连接右侧子通道
            if col < 5:  # 不是最右列
                subchannel.add_connection(
                    subchannel.id + 1,
                    'lateral',
                    self.dx
                )
            
            # 连接下方子通道
            if row < 5:  # 不是最底行
                subchannel.add_connection(
                    subchannel.id + 6,
                    'lateral',
                    self.dx
                )
    
    def initialize_field_variables(self):
        """初始化场变量"""
        for subchannel in self.subchannels:
            # 初始化空泡份额
            subchannel.alpha = np.zeros(self.n_axial)
            subchannel.previous_alpha = np.zeros(self.n_axial)
            
            # 初始化压力场
            subchannel.pressure = np.ones(self.n_axial) * 1.0e5  # 假设初始压力为1bar
            subchannel.previous_pressure = np.ones(self.n_axial) * 1.0e5
            
            # 初始化速度场
            for phase in ['liquid', 'vapor']:
                subchannel.velocity[phase]['axial'] = np.zeros(self.n_axial)
                subchannel.velocity[phase]['lateral'] = np.zeros(self.n_axial)
                subchannel.previous_velocity[phase]['axial'] = np.zeros(self.n_axial)
                subchannel.previous_velocity[phase]['lateral'] = np.zeros(self.n_axial)
                subchannel.temperature[phase] = np.ones(self.n_axial) * 300.0  # 假设初始温度为300K
                subchannel.density[phase] = np.zeros(self.n_axial)
    
    def set_boundary_conditions(self, inlet_conditions, outlet_conditions):
        """设置边界条件"""
        # 入口边界条件
        for subchannel in self.subchannels[:36]:  # 第一层的所有子通道
            # 设置速度
            for phase in ['liquid', 'vapor']:
                subchannel.velocity[phase]['axial'][0] = inlet_conditions[phase]['velocity']
                # 使用全局温度
                subchannel.temperature[phase][0] = inlet_conditions['temperature']
            
            # 设置压力和空泡份额
            subchannel.pressure[0] = inlet_conditions['pressure']
            subchannel.alpha[0] = inlet_conditions['alpha']
        
        # 出口边界条件
        for subchannel in self.subchannels[-36:]:  # 最后一层的所有子通道
            subchannel.pressure[-1] = outlet_conditions['pressure']
            
    def advance_time(self, dt):
        """推进时间步"""
        self.current_time += dt
        self.time_step_number += 1
        self.dt = dt
        
        # 保存当前状态为上一时间步的值
        for subchannel in self.subchannels:
            subchannel.save_current_state() 
    
    def _create_all_subchannels(self):
        """创建所有子通道"""
        # 清空现有子通道列表
        self.subchannels = []
        
        # 创建每一层的子通道
        for k in range(self.n_axial):
            for i in range(36):  # 6x6网格
                channel_id = k * 36 + i
                self.create_subchannel(channel_id) 
    
    def get_adjacent_channels(self, channel_id):
        """获取指定通道的相邻通道
        Args:
            channel_id: 子通道ID
        Returns:
            adjacent_channels: 相邻通道列表
        """
        # 根据通道类型确定相邻通道
        channel = self.subchannels[channel_id]
        adjacent_ids = []
        
        # 根据通道位置确定相邻通道
        if channel.type == 'corner':
            # 角通道有2个相邻通道
            if channel_id == 0:  # 左上角
                adjacent_ids = [1, 3]
            elif channel_id == 2:  # 右上角
                adjacent_ids = [1, 5]
            elif channel_id == 6:  # 左下角
                adjacent_ids = [3, 7]
            elif channel_id == 8:  # 右下角
                adjacent_ids = [5, 7]
        
        elif channel.type == 'edge':
            # 边通道有3个相邻通道
            if channel_id == 1:  # 上边
                adjacent_ids = [0, 2, 4]
            elif channel_id == 3:  # 左边
                adjacent_ids = [0, 6, 4]
            elif channel_id == 5:  # 右边
                adjacent_ids = [2, 8, 4]
            elif channel_id == 7:  # 下边
                adjacent_ids = [6, 8, 4]
        
        elif channel.type == 'center':
            # 中心通道有4个相邻通道
            adjacent_ids = [1, 3, 5, 7]
        
        return [self.subchannels[i] for i in adjacent_ids] 